import cv2
import robomaster
from robomaster import robot
from robomaster import vision
import time
from robomaster import gimbal

my_distance = []

def sub_data_handler(sub_info):
    global my_distance
    my_distance = sub_info
class MarkerInfo:

    def __init__(self, x, y, w, h, info):
        self._x = x
        self._y = y
        self._w = w
        self._h = h
        self._info = info

    @property
    def pt1(self):
        return int((self._x - self._w / 2) * 1280), int((self._y - self._h / 2) * 720)

    @property
    def pt2(self):
        return int((self._x + self._w / 2) * 1280), int((self._y + self._h / 2) * 720)

    @property
    def center(self):
        return int(self._x * 1280), int(self._y * 720)

    @property
    def text(self):
        return self._info


markers = []


def on_detect_marker(marker_info):
    number = len(marker_info)
    markers.clear()
    for i in range(0, number):
        x, y, w, h, info = marker_info[i]
        markers.append(MarkerInfo(x, y, w, h, info))
        # print("marker:{0} x:{1}, y:{2}, w:{3}, h:{4}".format(info, x, y, w, h))


if __name__ == '__main__':
    ep_robot = robot.Robot()

    ep_robot.initialize(conn_type="ap")
    ep_sensor = ep_robot.sensor
    ep_sensor.sub_distance(freq=5, callback=sub_data_handler)
    ep_vision = ep_robot.vision
    ep_camera = ep_robot.camera
    ep_chassis = ep_robot.chassis
    ep_gimbal = ep_robot.gimbal
    ep_gimbal.moveto(pitch=0, yaw=0).wait_for_completed()
    ep_robot.set_robot_mode(mode="gimbal")  # 底盘云台自由
    ep_camera.start_video_stream(display=False)
    result = ep_vision.sub_detect_info(name="marker", callback=on_detect_marker)

    time.sleep(3)

    while True:
        img = ep_camera.read_cv2_image(strategy="newest", timeout=2)
        for j in range(0, len(markers)):
            cv2.rectangle(img, markers[j].pt1, markers[j].pt2, (255, 255, 255))
            cv2.putText(img, markers[j].text, markers[j].center, cv2.FONT_HERSHEY_SIMPLEX, 1.5, (255, 255, 255), 3)
            #print(markers[0]._x)

        if len(markers) > 0:
            position_x = markers[0]._x

            # print(position_x)
            # print(my_distance[0])
            H =markers[0]._h
            error = position_x - 0.5
            output = 2 * error
            # print(output)
            ep_chassis.drive_speed(x=0, y=output, z=0)

            if my_distance[0] == 0:
                ep_chassis.drive_speed(x=0, y=0, z=0, timeout=5)
            else:
                if H < 0.2:  # 静止区
                    ep_chassis.drive_speed(x=0.4, y=output, z=0, timeout=3)
                elif 0.2 < H< 0.25:  # 跟进区
                    ep_chassis.drive_speed(x=0.2, y=output, z=0, timeout=1)
                elif H >0.25:  # 稳定区
                    ep_chassis.drive_speed(x=0, y=0, z=0)
                elif 450 >= my_distance[0] > 1:  # 后退区
                    ep_chassis.drive_speed(x=-0.2, y=output, z=0, timeout=3)
                else:
                    ep_chassis.drive_speed(x=0, y=0, z=0, timeout=1)
                print(my_distance[0])
                print(H)

        else:
            ep_chassis.drive_speed(x=0, y=0, z=0)

        cv2.imshow("Markers", img)
        key = cv2.waitKey(1)
        if key == ord('q'):
            break

    result = ep_vision.unsub_detect_info(name="marker")
    cv2.destroyAllWindows()
    ep_camera.stop_video_stream()
    ep_robot.close()
